import torch
import torch.nn as nn
import torch.nn.functional as F
import time, pdb, numpy

class angleproto(nn.Module):
    def __init__(self, w=10.0, b=-5.0, **kwargs):
        super(angleproto, self).__init__()
        self.w = nn.Parameter(torch.tensor(w))
        self.b = nn.Parameter(torch.tensor(b))
        self.criterion  = torch.nn.CrossEntropyLoss()
        print('Initialised AngleProto')

    def forward(self, out_anchor, out_positive):
        stepsize = out_anchor.size()[0]
        cos_sim_matrix  = F.cosine_similarity(out_positive.unsqueeze(-1),out_anchor.unsqueeze(-1).transpose(0,2))
        torch.clamp(self.w, 1e-6)
        cos_sim_matrix = cos_sim_matrix * self.w + self.b
        
        label = torch.from_numpy(numpy.asarray(range(0,stepsize))).cuda()
        loss = self.criterion(cos_sim_matrix, label)

        return loss
